#include <iostream>
#include <vector>
#include <queue>
#include <unordered_map>
#include <unordered_set>
#include <Eigen/Dense>
#include <Eigen/Geometry>

using namespace Eigen;
using namespace std;

const double PI = 3.14159265358979323846;

// 2D位姿 (x, y, theta)
struct Pose2D {
    double x, y, theta;
    Pose2D(double x = 0, double y = 0, double theta = 0) : x(x), y(y), theta(theta) {}
    
    Matrix3d toMatrix() const {
        Matrix3d T = Matrix3d::Identity();
        T(0, 0) = cos(theta); T(0, 1) = -sin(theta);
        T(1, 0) = sin(theta); T(1, 1) = cos(theta);
        T(0, 2) = x;
        T(1, 2) = y;
        return T;
    }
    
    static Pose2D fromMatrix(const Matrix3d& T) {
        return Pose2D(T(0, 2), T(1, 2), atan2(T(1, 0), T(0, 0)));
    }
};

// 位姿图中的边 (约束)
struct PoseGraphEdge {
    int id1, id2;           // 连接的顶点ID
    Pose2D measurement;     // 相对位姿测量
    Matrix3d information;   // 信息矩阵
    
    PoseGraphEdge(int id1, int id2, const Pose2D& meas, const Matrix3d& info)
        : id1(id1), id2(id2), measurement(meas), information(info) {}
};

// 滑动窗口位姿图优化器
class SlidingWindowPoseGraph {
public:
    SlidingWindowPoseGraph(int window_size = 10) : window_size(window_size) {}
    
    // 添加新顶点 (返回顶点ID)
    int addVertex(const Pose2D& pose, bool is_fixed = false) {
        int id = next_vertex_id++;
        vertices[id] = pose;
        isFixed[id] = is_fixed;
        vertex_queue.push(id);
        
        // 如果超过窗口大小，边缘化最老的顶点
        if (vertex_queue.size() > window_size) {
            marginalizeOldestVertex();
        }
        
        return id;
    }
    
    // 添加边约束
    void addEdge(const PoseGraphEdge& edge) {
        edges.push_back(edge);
        size_t edge_idx = edges.size() - 1;
        active_edges.insert(edge_idx);
        vertex_to_edges[edge.id1].insert(edge_idx);
        vertex_to_edges[edge.id2].insert(edge_idx);
    }
    
    // 设置先验
    void setPrior(int vertex_id, const Pose2D& prior_pose, const Matrix3d& prior_info) {
        priors[vertex_id] = {prior_pose, prior_info};
    }
    
    // 执行优化
    void optimize(int max_iterations = 5) {
        for (int iter = 0; iter < max_iterations; ++iter) {
            buildSystem();
            
            // 添加阻尼因子确保正定性
            for (int i = 0; i < H.rows(); ++i) {
                H(i, i) += 1e-6;
            }
            
            VectorXd dx = H.ldlt().solve(-b);
            updatePoses(dx);
        }
    }
    
    // 获取位姿
    Pose2D getPose(int id) const {
        auto it = vertices.find(id);
        if (it == vertices.end()) throw runtime_error("Vertex not found");
        return it->second;
    }
    
    // 打印当前位姿
    void printPoses() const {
        cout << "Current poses in window:" << endl;
        for (const auto& v : vertices) {
            cout << "Vertex " << v.first << ": (" << v.second.x << ", " 
                 << v.second.y << ", " << v.second.theta << ")" << endl;
        }
    }

private:
    int window_size;
    int next_vertex_id = 0;
    queue<int> vertex_queue;
    unordered_map<int, Pose2D> vertices;
    unordered_map<int, bool> isFixed;
    vector<PoseGraphEdge> edges;
    unordered_set<size_t> active_edges;
    unordered_map<int, pair<Pose2D, Matrix3d>> priors;
    unordered_map<int, unordered_set<size_t>> vertex_to_edges;
    
    // 边缘化因子
    struct MarginalizationFactor {
        MatrixXd H;
        VectorXd b;
        unordered_set<int> related_vertices;
    };
    unordered_map<int, MarginalizationFactor> marginalization_factors;
    
    // 优化系统
    MatrixXd H;
    VectorXd b;
    unordered_map<int, int> var_indices;
    
    // 边缘化最老的顶点
    void marginalizeOldestVertex() {
        if (vertex_queue.empty()) return;
        
        int marginalize_id = vertex_queue.front();
        vertex_queue.pop();
        
        if (isFixed[marginalize_id]) {
            vertices.erase(marginalize_id);
            isFixed.erase(marginalize_id);
            vertex_to_edges.erase(marginalize_id);
            return;
        }
        
        buildSystemForMarginalization(marginalize_id);

        MatrixXd H_marg;
        VectorXd b_marg;
        std::tie(H_marg, b_marg) = performSchurComplement(marginalize_id);
        
        // 存储边缘化因子
        MarginalizationFactor factor;
        factor.H = H_marg;
        factor.b = b_marg;
        
        // 找出相关顶点
        for (size_t edge_idx : vertex_to_edges[marginalize_id]) {
            const auto& edge = edges[edge_idx];
            if (edge.id1 != marginalize_id) factor.related_vertices.insert(edge.id1);
            if (edge.id2 != marginalize_id) factor.related_vertices.insert(edge.id2);
        }
        
        // 添加先验相关的顶点
        for (const auto& prior : priors) {
            if (prior.first != marginalize_id && vertices.count(prior.first)) {
                factor.related_vertices.insert(prior.first);
            }
        }
        
        marginalization_factors[marginalize_id] = factor;
        
        // 清理
        vertices.erase(marginalize_id);
        isFixed.erase(marginalize_id);
        vertex_to_edges.erase(marginalize_id);
    }
    
    // 为边缘化构建子系统
    void buildSystemForMarginalization(int marginalize_id) {
        // 找出相关边
        vector<size_t> related_edges;
        for (size_t i = 0; i < edges.size(); ++i) {
            if (active_edges.count(i) && 
                (edges[i].id1 == marginalize_id || edges[i].id2 == marginalize_id)) {
                related_edges.push_back(i);
            }
        }
        
        // 构建局部Hessian和右侧
        int n = vertices.size();
        H = MatrixXd::Zero(3*n, 3*n);
        b = VectorXd::Zero(3*n);
        
        // 创建变量索引映射
        var_indices.clear();
        int idx = 0;
        for (const auto& v : vertices) {
            var_indices[v.first] = idx++;
        }
        
        // 添加相关边约束
        for (size_t edge_idx : related_edges) {
            const auto& edge = edges[edge_idx];
            int i = edge.id1, j = edge.id2;
            
            Pose2D xi = vertices[i], xj = vertices[j];
            Pose2D z = edge.measurement;
            Matrix3d Omega = edge.information;
               
            // 计算误差和雅可比
            Matrix2d R_i = Rotation2Dd(xi.theta).toRotationMatrix();
            Matrix2d R_z = Rotation2Dd(z.theta).toRotationMatrix();

            Vector3d e;
            e(0) = R_z(0,0)*(xj.x-xi.x) + R_z(1,0)*(xj.y-xi.y) - z.x;
            e(1) = -R_z(0,1)*(xj.x-xi.x) + R_z(1,1)*(xj.y-xi.y) - z.y;
            e(2) = normalizeAngle(xj.theta - xi.theta - z.theta);

            Matrix3d Ji = Matrix3d::Zero();
            Ji.block<2,2>(0,0) = -R_z.transpose() * R_i.transpose();
            Ji(0,2) = -R_i(1,0)*(xj.x-xi.x) + R_i(1,1)*(xj.y-xi.y);
            Ji(1,2) = R_i(0,0)*(xj.x-xi.x) - R_i(0,1)*(xj.y-xi.y);
            Ji(2,2) = -1;

            Matrix3d Jj = Matrix3d::Zero();
            Jj.block<2,2>(0,0) = R_z.transpose() * R_i.transpose();
            Jj(2,2) = 1;

            // 更新局部系统
            int idx_i = var_indices[i], idx_j = var_indices[j];
            
            H.block<3, 3>(3*idx_i, 3*idx_i) += Ji.transpose() * Omega * Ji;
            H.block<3, 3>(3*idx_i, 3*idx_j) += Ji.transpose() * Omega * Jj;
            H.block<3, 3>(3*idx_j, 3*idx_i) += Jj.transpose() * Omega * Ji;
            H.block<3, 3>(3*idx_j, 3*idx_j) += Jj.transpose() * Omega * Jj;
            
            b.segment<3>(3*idx_i) += Ji.transpose() * Omega * e;
            b.segment<3>(3*idx_j) += Jj.transpose() * Omega * e;
        }
        
        // 添加先验约束
        for (const auto& prior : priors) {
            int i = prior.first;
            if (vertices.count(i) == 0) continue;
            
            int idx_i = var_indices[i];
            Pose2D prior_pose = prior.second.first;
            Matrix3d prior_info = prior.second.second;
            
            Pose2D xi = vertices[i];
            Vector3d e;
            e << xi.x - prior_pose.x, xi.y - prior_pose.y, normalizeAngle(xi.theta - prior_pose.theta);
            
            H.block<3, 3>(3*idx_i, 3*idx_i) += prior_info;
            b.segment<3>(3*idx_i) += prior_info * e;
        }
    }
    
    // 执行舒尔补边缘化
    pair<MatrixXd, VectorXd> performSchurComplement(int marginalize_id) {
        int marg_idx = var_indices[marginalize_id];
        
        // 分割变量: α=要边缘化的变量, β=保留的变量
        int n = var_indices.size();
        vector<int> keep_indices;
        for (int i = 0; i < n; ++i) {
            if (i != marg_idx) keep_indices.push_back(i);
        }
        
        // 提取矩阵块
        MatrixXd H_alpha_alpha = H.block<3, 3>(3*marg_idx, 3*marg_idx);
        MatrixXd H_alpha_beta = H.block(3*marg_idx, 0, 3, 3*(n-1));
        MatrixXd H_beta_alpha = H.block(0, 3*marg_idx, 3*(n-1), 3);
        MatrixXd H_beta_beta = MatrixXd::Zero(3*(n-1), 3*(n-1));
        
        VectorXd b_alpha = b.segment<3>(3*marg_idx);
        VectorXd b_beta = VectorXd::Zero(3*(n-1));
        
        // 构建保留部分的H_beta_beta
        for (size_t i = 0; i < keep_indices.size(); ++i) {
            for (size_t j = 0; j < keep_indices.size(); ++j) {
                int row = keep_indices[i], col = keep_indices[j];
                H_beta_beta.block<3, 3>(3*i, 3*j) = H.block<3, 3>(3*row, 3*col);
            }
            b_beta.segment<3>(3*i) = b.segment<3>(3*keep_indices[i]);
        }
        
        // 计算舒尔补
        LLT<MatrixXd> llt_Haa(H_alpha_alpha);
        MatrixXd H_aa_inv_H_ab = llt_Haa.solve(H_alpha_beta);
        
        MatrixXd H_marg = H_beta_beta - H_beta_alpha * H_aa_inv_H_ab;
        VectorXd b_marg = b_beta - H_beta_alpha * llt_Haa.solve(b_alpha);
        
        return {H_marg, b_marg};
    }
    
    // 构建全局优化系统
    void buildSystem() {
        // 确定自由变量
        var_indices.clear();
        int free_vars = 0;
        for (const auto& v : vertices) {
            if (!isFixed[v.first]) {
                var_indices[v.first] = free_vars++;
            }
        }
        
        // 初始化Hessian和右侧
        H = MatrixXd::Zero(3 * free_vars, 3 * free_vars);
        b = VectorXd::Zero(3 * free_vars);
        
        // 添加边约束
        for (size_t i = 0; i < edges.size(); ++i) {
            if (active_edges.count(i) == 0) continue;
            
            const auto& edge = edges[i];
            int id1 = edge.id1, id2 = edge.id2;
            bool id1_fixed = isFixed[id1], id2_fixed = isFixed[id2];
            
            if (id1_fixed && id2_fixed) continue;
            
            Pose2D xi = vertices[id1], xj = vertices[id2];
            Pose2D z = edge.measurement;
            Matrix3d Omega = edge.information;
                      
            // 计算误差和雅可比
            Matrix2d R_i = Rotation2Dd(xi.theta).toRotationMatrix();
            Matrix2d R_z = Rotation2Dd(z.theta).toRotationMatrix();

            Vector3d e;
            e(0) = R_z(0,0)*(xj.x-xi.x) + R_z(1,0)*(xj.y-xi.y) - z.x;
            e(1) = -R_z(0,1)*(xj.x-xi.x) + R_z(1,1)*(xj.y-xi.y) - z.y;
            e(2) = normalizeAngle(xj.theta - xi.theta - z.theta);

            Matrix3d Ji = Matrix3d::Zero();
            Ji.block<2,2>(0,0) = -R_z.transpose() * R_i.transpose();
            Ji(0,2) = -R_i(1,0)*(xj.x-xi.x) + R_i(1,1)*(xj.y-xi.y);
            Ji(1,2) = R_i(0,0)*(xj.x-xi.x) - R_i(0,1)*(xj.y-xi.y);
            Ji(2,2) = -1;

            Matrix3d Jj = Matrix3d::Zero();
            Jj.block<2,2>(0,0) = R_z.transpose() * R_i.transpose();
            Jj(2,2) = 1;

            // 更新系统
            if (!id1_fixed && !id2_fixed) {
                int idx1 = var_indices[id1], idx2 = var_indices[id2];
                
                H.block<3, 3>(3*idx1, 3*idx1) += Ji.transpose() * Omega * Ji;
                H.block<3, 3>(3*idx1, 3*idx2) += Ji.transpose() * Omega * Jj;
                H.block<3, 3>(3*idx2, 3*idx1) += Jj.transpose() * Omega * Ji;
                H.block<3, 3>(3*idx2, 3*idx2) += Jj.transpose() * Omega * Jj;
                
                b.segment<3>(3*idx1) += Ji.transpose() * Omega * e;
                b.segment<3>(3*idx2) += Jj.transpose() * Omega * e;
            } else if (!id1_fixed) {
                int idx1 = var_indices[id1];
                H.block<3, 3>(3*idx1, 3*idx1) += Ji.transpose() * Omega * Ji;
                b.segment<3>(3*idx1) += Ji.transpose() * Omega * e;
            } else if (!id2_fixed) {
                int idx2 = var_indices[id2];
                H.block<3, 3>(3*idx2, 3*idx2) += Jj.transpose() * Omega * Jj;
                b.segment<3>(3*idx2) += Jj.transpose() * Omega * e;
            }
        }
        
        // 添加先验约束
        for (const auto& prior : priors) {
            int id = prior.first;
            if (isFixed[id] || vertices.count(id) == 0) continue;
            
            int idx = var_indices[id];
            Pose2D prior_pose = prior.second.first;
            Matrix3d prior_info = prior.second.second;
            
            Pose2D xi = vertices[id];
            Vector3d e;
            e << xi.x - prior_pose.x, xi.y - prior_pose.y, normalizeAngle(xi.theta - prior_pose.theta);
            
            H.block<3, 3>(3*idx, 3*idx) += prior_info;
            b.segment<3>(3*idx) += prior_info * e;
        }
        
        // 添加边缘化因子
        for (const auto& factor_pair : marginalization_factors) {
            const MarginalizationFactor factor = factor_pair.second;
            
            // 找出当前窗口中存在的相关顶点
            vector<int> active_related_vertices;
            for (int vid : factor.related_vertices) {
                if (vertices.count(vid)) {
                    active_related_vertices.push_back(vid);
                }
            }
            
            if (active_related_vertices.empty()) continue;
            
            // 创建局部变量索引映射
            unordered_map<int, int> local_indices;
            int local_idx = 0;
            for (int vid : active_related_vertices) {
                if (!isFixed[vid]) {
                    local_indices[vid] = local_idx++;
                }
            }
            int num_local_vars = local_idx;
            
            if (num_local_vars == 0) continue;
            
            // 提取相关的H和b块
            MatrixXd H_factor = MatrixXd::Zero(3*num_local_vars, 3*num_local_vars);
            VectorXd b_factor = VectorXd::Zero(3*num_local_vars);
            
            // 构建从全局变量到局部变量的映射
            vector<int> global_indices;
            for (int vid : active_related_vertices) {
                if (!isFixed[vid]) {
                    global_indices.push_back(var_indices[vid]);
                }
            }
            
            // 填充H_factor和b_factor
            for (int i = 0; i < num_local_vars; ++i) {
                for (int j = 0; j < num_local_vars; ++j) {
                    H_factor.block<3,3>(3*i, 3*j) = 
                        factor.H.block<3,3>(3*global_indices[i], 3*global_indices[j]);
                }
                b_factor.segment<3>(3*i) = factor.b.segment<3>(3*global_indices[i]);
            }
            
            // 将因子添加到全局系统
            for (int i = 0; i < num_local_vars; ++i) {
                int global_i = global_indices[i];
                for (int j = 0; j < num_local_vars; ++j) {
                    int global_j = global_indices[j];
                    H.block<3,3>(3*global_i, 3*global_j) += H_factor.block<3,3>(3*i, 3*j);
                }
                b.segment<3>(3*global_i) += b_factor.segment<3>(3*i);
            }
        }
    }
    
    // 更新位姿
    void updatePoses(const VectorXd& dx) {
        for (const auto& v : var_indices) {
            int id = v.first;
            int idx = v.second;
            vertices[id].x += dx(3*idx);
            vertices[id].y += dx(3*idx + 1);
            vertices[id].theta = normalizeAngle(vertices[id].theta + dx(3*idx + 2));
        }
    }
    
    // 角度归一化
    static double normalizeAngle(double angle) {
        while (angle > PI) angle -= 2 * PI;
        while (angle < -PI) angle += 2 * PI;
        return angle;
    }
};

int main() {
    // 创建滑动窗口优化器 (窗口大小=5)
    SlidingWindowPoseGraph optimizer(5);
    
    // 添加初始顶点 (固定)
    int id0 = optimizer.addVertex(Pose2D(0, 0, 0), true);
    
    // 模拟机器人运动轨迹
    vector<int> ids;
    for (int i = 1; i <= 10; ++i) {
        Pose2D pose(i * 1.0, 0, 0); // 直线运动
        int id = optimizer.addVertex(pose);
        ids.push_back(id);
        
        // 添加里程计边
        if (i > 1) {
            Matrix3d info_matrix  = Matrix3d::Identity();
            info_matrix(2,2) = 100.0;  // 给角度更高的权重
            optimizer.addEdge(PoseGraphEdge(ids[i-2], ids[i-1], Pose2D(1.0, 0, 0), info_matrix ));

            // 添加运动约束
            Matrix3d motion_info = Matrix3d::Identity();
            motion_info(0,0) = 10000.0;
            motion_info(1,1) = 1000.0;
            motion_info(2,2) = 1000.0;
            optimizer.addEdge(PoseGraphEdge(ids[i-2], ids[i-1], Pose2D(1.0, 0, 0), motion_info));
        }
        
        // 添加回环边 (每5步)
        if (i % 5 == 0) {
            Matrix3d loop_info = 100000 * Matrix3d::Identity();
            loop_info(2,2) = 1000.0;
            optimizer.addEdge(PoseGraphEdge(id0, ids[i-1], Pose2D(i * 1.0, 0, 0), loop_info));
        }
        
        // 执行优化
        if (i >= 2) {
            optimizer.optimize(10);
            cout << "After step " << i << ":" << endl;
            optimizer.printPoses();
            cout << "----------------------" << endl;
        }
    }
    
    // 输出最终位姿
    cout << "\nFinal optimized poses:" << endl;
    for (size_t i = 0; i < ids.size(); ++i) {
        Pose2D pose = optimizer.getPose(ids[i]);
        cout << "Vertex " << ids[i] << ": (" << pose.x << ", " << pose.y << ", " << pose.theta << ")" << endl;
    }
    
    return 0;
}